#!/usr/bin/env python3

from subprocess import call
from sys import argv, stderr, stdout, exit
from os.path import isfile, join, dirname, basename, abspath
from os import chdir, getcwd
from tempfile import mkdtemp
from shutil import rmtree

wd = None
olddir = getcwd()
debug = False

def err(msg):
    stdout.flush()
    stderr.flush()
    stderr.write('[dgtool]: ' + msg + '\n')

    if wd:
        chdir(olddir);
        printcmd(["rmtree", wd])
        rmtree(wd)
    exit(1)

def printcmd(cmd):
    if debug:
        print('> {0}'.format(' '.join(cmd)))

prog=None
clang='clang'
link='llvm-link'

clangcmd=[clang, '-emit-llvm', '-c', '-g']
cmd=[]
files=[]
bcfiles=[]
yamlfiles=[]

def get_tool(prog):
    # try to use the binaries from the tools/ directory
    path=f"./{prog}"
    if isfile(path):
        prog = abspath(path)
    else:
        path = join(dirname(argv[0]), prog)
        if isfile(path):
            prog = abspath(path)
        # else keep prog untouched,
        # maybe it is in PATH var
    return prog

###
# Build the commands
options=[]
passes=[]
next_is_dg=False
next_is_clang=False
next_is_opt=False
all_is_clang=False
cc_mode = False
for x in argv[1:]:
    if next_is_clang:
        clangcmd.append(x)
        next_is_clang = False
    elif next_is_dg:
        options.append(x)
        next_is_dg = False
    elif next_is_opt:
        passes.append(x)
        next_is_opt = False
    elif x == '-Xclang':
        next_is_clang = True
    elif x == '-Xdg':
        next_is_dg = True
    elif x == '-Xopt':
        next_is_opt = True
    elif x.endswith('.c') or x.endswith('.cpp') or x.endswith('.i'):
        files.append(abspath(x))
    elif x.endswith('.bc') or x.endswith('.ll'):
        bcfiles.append(abspath(x))
    elif x.endswith('.yml') or x.endswith('.yaml'):
        yamlfiles.append(abspath(x))
    else:
        if all_is_clang: # (all except files)
            clangcmd.append(x)
        elif not cmd:
            # cc mode we just compile, so everything goes to clang
            if x == 'cc':
                all_is_clang = True
                cc_mode = True
                continue
            cmd.append(get_tool(x))
        else:
            cmd.append(x)

if 'debug' in options or 'dbg' in options:
    debug=True
###
# Try to find the right path to the program
if not cmd and not cc_mode:
    err('no command given')

if yamlfiles:
    try:
        import yaml
    except ImportError:
        err("Got YAML file, but does not have yaml module")

    for fl in yamlfiles:
        with open(fl) as yf:
            data = yaml.load(yf, Loader=yaml.FullLoader)
            inp = data['input_files'] #NOTE is always just one?
            files.append(join(dirname(fl), inp))

def repl_suffix(fl):
    f = basename(fl)
    if f.endswith('.c') or f.endswith('.i'):
        return '{0}.bc'.format(f[:-2])
    elif f.endswith('.cpp'):
        return '{0}.bc'.format(f[:-4])

if len(files) < 1 and len(bcfiles) < 1:
    err('No input files given')

###
# Create temporary directory
if not cc_mode:
    wd = mkdtemp(suffix='dgtool')
    printcmd(["chdir", wd])
    chdir(wd)

###
# Issue clang to compile sources to bitcode(s)
clangcmd += files
bitcodes = bcfiles + list(map(repl_suffix, files))
bitcode = None

printcmd(clangcmd)
r = call(clangcmd)
if r != 0:
    err('clang command return non-zero status')

###
# Link multiple files if needed
if len(bitcodes) > 1:
    bitcode='bitcode.bc'
    linkcmd=[link, '-o', bitcode] + bitcodes

    printcmd(linkcmd)
    r = call(linkcmd)
    if r != 0:
        err('llvm-link command return non-zero status')

else:
    bitcode = bitcodes[0]

if passes:
    newname = bitcode + '.opt.bc'
    optcmd = ['opt', bitcode, '-o', newname] + passes
    printcmd(optcmd)
    r = call(optcmd)
    if r != 0:
        err('opt command return non-zero status')

    bitcode = newname

###
# Run the tool
if not cc_mode:
    cmd.append(bitcode)
    printcmd(cmd)

    try:
        r = call(cmd)
        if r != 0:
            err('command return non-zero status')
    except OSError as e:
            err('command failed: {0}'.format(str(e)))

if wd:
    chdir(olddir);
    printcmd(["rmtree", wd])
    rmtree(wd)
exit(0)
